{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import backend as K\n",
    "\n",
    "class PositionalEncoding(tf.keras.layers.Layer):\n",
    "    def __init__(self, sequence_len=None, embedding_dim=None, mode='sum', **kwargs):\n",
    "        self.sequence_len = sequence_len\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.mode = mode\n",
    "        super(PositionalEncoding, self).__init__(**kwargs)\n",
    "\n",
    "    def call(self, x):\n",
    "        if (self.embedding_dim == None) or (self.mode == 'sum'):\n",
    "            self.embedding_dim = int(x.shape[-1])\n",
    "        \n",
    "        position_embedding = np.array([\n",
    "            [pos / np.power(10000, 2. * i / self.embedding_dim) for i in range(self.embedding_dim)]\n",
    "            for pos in range(self.sequence_len)])\n",
    "\n",
    "        position_embedding[:, 0::2] = np.sin(position_embedding[:, 0::2])  # dim 2i\n",
    "        position_embedding[:, 1::2] = np.cos(position_embedding[:, 1::2])  # dim 2i+1\n",
    "        \n",
    "        position_embedding = tf.cast(position_embedding, dtype=tf.float32)\n",
    "        \n",
    "        if self.mode == 'sum':\n",
    "            return position_embedding + x\n",
    "        \n",
    "        elif self.mode == 'concat':\n",
    "            position_embedding = tf.reshape(\n",
    "                tf.tile(position_embedding, (int(x.shape[0]), 1)), \n",
    "                (-1, self.sequence_len, self.embedding_dim)\n",
    "                )\n",
    "\n",
    "            return tf.concat([position_embedding, x], 2)\n",
    "        \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        if self.mode == 'sum':\n",
    "            return input_shape\n",
    "        \n",
    "        elif self.mode == 'concat':\n",
    "            return (input_shape[0], input_shape[1], input_shape[2]+self.embedding_dim)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[ 0.09794044  1.5221956   0.05093884 ...  1.2205979   0.6255225\n",
      "    1.9426765 ]\n",
      "  [ 1.2793872   1.707086    1.4439819  ...  1.9408028   0.14384009\n",
      "    1.64903   ]\n",
      "  [ 1.4325817   0.57568574  1.7065923  ...  1.8288414   0.07602505\n",
      "    1.646029  ]\n",
      "  ...\n",
      "  [ 0.7831629   0.20508784  1.539695   ...  1.2742875   0.90087354\n",
      "    1.5430839 ]\n",
      "  [-0.11898029  0.08894062  1.2983346  ...  1.5916389   0.9796848\n",
      "    1.623994  ]\n",
      "  [-0.38765883  0.93811846  1.0189443  ...  1.5155734   0.62063557\n",
      "    1.49345   ]]\n",
      "\n",
      " [[ 0.75529814  1.7238648   0.30029643 ...  1.4552534   0.05495417\n",
      "    1.0411552 ]\n",
      "  [ 1.2798775   1.2864672   0.93499756 ...  1.2959702   0.0053605\n",
      "    1.846127  ]\n",
      "  [ 1.7781134   0.8382853   1.0850062  ...  1.5704403   0.604845\n",
      "    1.371224  ]\n",
      "  ...\n",
      "  [ 0.32015985 -0.54433995  1.4011699  ...  1.1479545   0.7128234\n",
      "    1.8522563 ]\n",
      "  [-0.6383225   0.08011305  1.6697282  ...  1.1402829   0.5857736\n",
      "    1.4564639 ]\n",
      "  [-0.34154665  0.9058945   0.6852306  ...  1.6051351   0.24876395\n",
      "    1.3167462 ]]\n",
      "\n",
      " [[ 0.07864439  1.3381037   0.5568241  ...  1.057088    0.9738351\n",
      "    1.1073668 ]\n",
      "  [ 1.6080778   0.8790256   1.1625597  ...  1.8691885   0.7797855\n",
      "    1.009449  ]\n",
      "  [ 1.6426668   0.2203561   1.486139   ...  1.7387081   0.40477636\n",
      "    1.1706511 ]\n",
      "  ...\n",
      "  [ 0.97335654 -0.50888664  1.2674762  ...  1.7055811   0.79238236\n",
      "    1.1494101 ]\n",
      "  [-0.6044022   0.18841112  1.8560048  ...  1.2342762   0.90584743\n",
      "    1.3253067 ]\n",
      "  [-0.6953174   1.2022023   0.74448514 ...  1.5259538   0.5035129\n",
      "    1.3181292 ]]\n",
      "\n",
      " ...\n",
      "\n",
      " [[ 0.9582726   1.3343097   0.5962509  ...  1.1495905   0.11049569\n",
      "    1.0063419 ]\n",
      "  [ 1.6500767   1.5479307   0.568589   ...  1.1155576   0.45289448\n",
      "    1.0962205 ]\n",
      "  [ 1.2926505   0.8690175   1.3382267  ...  1.0641884   0.49016264\n",
      "    1.0090787 ]\n",
      "  ...\n",
      "  [ 0.7488863  -0.311369    1.6898346  ...  1.8387809   0.63505507\n",
      "    1.6838411 ]\n",
      "  [-0.02913785  0.8243176   1.6379929  ...  1.6559383   0.69062364\n",
      "    1.6334406 ]\n",
      "  [-0.64841676  1.2795874   0.6969944  ...  1.9289603   0.28751984\n",
      "    1.4255868 ]]\n",
      "\n",
      " [[ 0.5696068   1.0768652   0.05855155 ...  1.7072641   0.22183394\n",
      "    1.0093197 ]\n",
      "  [ 1.4745699   0.988994    0.7231319  ...  1.4183924   0.12126292\n",
      "    1.9231985 ]\n",
      "  [ 1.2409298   0.20401394  1.7807915  ...  1.66094     0.9816645\n",
      "    1.825705  ]\n",
      "  ...\n",
      "  [ 0.8658504  -0.28001982  1.0415051  ...  1.8331397   0.73066413\n",
      "    1.371026  ]\n",
      "  [-0.22472656  0.40741813  1.2816292  ...  1.6981409   0.57380354\n",
      "    1.0430838 ]\n",
      "  [-0.6337311   0.6165538   0.978809   ...  1.4059229   0.27378252\n",
      "    1.9214246 ]]\n",
      "\n",
      " [[ 0.07307875  1.7860353   0.6352079  ...  1.8677078   0.06782842\n",
      "    1.7553717 ]\n",
      "  [ 0.99217236  1.3695343   1.2698109  ...  1.2379968   0.89337695\n",
      "    1.1643444 ]\n",
      "  [ 1.1061301   0.89300406  1.0930624  ...  1.4132322   0.52756995\n",
      "    1.1292539 ]\n",
      "  ...\n",
      "  [ 0.22643775 -0.25445193  1.0119381  ...  1.9650794   0.30556905\n",
      "    1.2038081 ]\n",
      "  [-0.62414706  0.76371694  1.4416256  ...  1.4597278   0.02213516\n",
      "    1.3326037 ]\n",
      "  [-0.64550316  1.1963696   1.1019957  ...  1.9422288   0.35251227\n",
      "    1.4477086 ]]], shape=(10, 50, 64), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# PositionalEncoding 测试\n",
    "\n",
    "position_embedding_layer = PositionalEncoding(50, 64, 'sum')\n",
    "position_embedding_layer_output = position_embedding_layer(tf.random.uniform((10, 50, 64)))\n",
    "\n",
    "print(position_embedding_layer_output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "def padding_mask(seq):\n",
    "    \n",
    "    # 获取为 0的padding项\n",
    "    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
    "\n",
    "    # 扩充维度用于attention矩阵\n",
    "    return seq[:, np.newaxis, np.newaxis, :] # (batch_size, 1, 1, seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[[1. 0. 0. 0.]]]\n",
      "\n",
      "\n",
      " [[[0. 0. 0. 1.]]]], shape=(2, 1, 1, 4), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# padding mask 测试\n",
    "\n",
    "padding_mask_list = padding_mask([[0,1,2,3], \n",
    "              [3,4,5,0]])\n",
    "\n",
    "print(padding_mask_list)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(q, k, v, mask):\n",
    "    \n",
    "    matmul_qk = tf.matmul(q, k, transpose_b=True)\n",
    "    dim_k = tf.cast(tf.shape(k)[-1], tf.float32)\n",
    "    scaled_attention_logits = matmul_qk / tf.math.sqrt(dim_k)\n",
    "    \n",
    "    if mask is not None:\n",
    "        scaled_attention_logits += (mask * -1e9)\n",
    "\n",
    "    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)\n",
    "    output = tf.matmul(attention_weights, v)\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "test_out: \n",
      "tf.Tensor(\n",
      "[[[1.2379461e-17 1.0000000e+01 7.0140370e-17 2.0000000e+00 1.0000000e+00\n",
      "   1.0000000e+00]]\n",
      "\n",
      " [[1.2375409e-07 2.9999888e+00 9.9998913e+00 1.0000000e+00 9.9999988e-01\n",
      "   1.0000000e+00]]], shape=(2, 1, 6), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# dot_product_attention 测试\n",
    "\n",
    "test_k = tf.constant([[[10,0,0,1,1,1],\n",
    "                      [0,10,0,2,1,1],\n",
    "                      [0,0,10,3,1,1],\n",
    "                      [0,0,10,4,1,1]],\n",
    "                      \n",
    "                      [[1,0,0,0,1,1],\n",
    "                      [0,2,0,1,1,1],\n",
    "                      [0,3,10,1,1,1],\n",
    "                      [1,0,4,1,0,1]]], dtype=tf.float32)  # (2, 4, 6)\n",
    "\n",
    "test_v = tf.constant([[[10,0,0,1,1,1],\n",
    "                      [0,10,0,2,1,1],\n",
    "                      [0,0,10,3,1,1],\n",
    "                      [0,0,10,4,1,1]],\n",
    "                      \n",
    "                      [[1,0,0,0,1,1],\n",
    "                      [0,2,0,1,1,1],\n",
    "                      [0,3,10,1,1,1],\n",
    "                      [1,0,4,1,0,1]]], dtype=tf.float32)  # (2, 4, 6)\n",
    "\n",
    "test_q = tf.constant([[[0,10,0,1,1,1]],\n",
    "                      [[0,8,2,5,3,1]]], dtype=tf.float32)\n",
    "\n",
    "\n",
    "mask = None\n",
    "test_out = scaled_dot_product_attention(test_q, test_k, test_v, mask)\n",
    "\n",
    "#mask = padding_mask(tf.constant([[1, 2, 0, 0], [3, 0, 1, 1]]))\n",
    "#test_out = scaled_dot_product_attention(test_k, test_k, test_v, mask)\n",
    "\n",
    "print(\"\\n\")\n",
    "print(\"test_out: \")\n",
    "print(test_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "querys: \n",
      "tf.Tensor(\n",
      "[[[[ 0. 10.  0.]]\n",
      "\n",
      "  [[ 1.  1.  1.]]]\n",
      "\n",
      "\n",
      " [[[ 0.  8.  2.]]\n",
      "\n",
      "  [[ 5.  3.  1.]]]], shape=(2, 2, 1, 3), dtype=float32)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "multi_head_test_out: \n",
      "tf.Tensor(\n",
      "[[[[8.4332741e-25 1.0000000e+01 0.0000000e+00]]\n",
      "\n",
      "  [[1.6404574e+00 1.0000000e+00 1.0000000e+00]]]\n",
      "\n",
      "\n",
      " [[[9.4977648e-10 3.0000000e+00 1.0000000e+01]]\n",
      "\n",
      "  [[9.5476758e-01 8.5647416e-01 1.0000000e+00]]]], shape=(2, 2, 1, 3), dtype=float32)\n",
      "\n",
      "\n",
      "scaled_attention: \n",
      "tf.Tensor(\n",
      "[[[[8.4332741e-25 1.0000000e+01 0.0000000e+00]\n",
      "   [1.6404574e+00 1.0000000e+00 1.0000000e+00]]]\n",
      "\n",
      "\n",
      " [[[9.4977648e-10 3.0000000e+00 1.0000000e+01]\n",
      "   [9.5476758e-01 8.5647416e-01 1.0000000e+00]]]], shape=(2, 1, 2, 3), dtype=float32)\n",
      "\n",
      "\n",
      "concat_attention: \n",
      "tf.Tensor(\n",
      "[[[8.4332741e-25 1.0000000e+01 0.0000000e+00 1.6404574e+00 1.0000000e+00\n",
      "   1.0000000e+00]]\n",
      "\n",
      " [[9.4977648e-10 3.0000000e+00 1.0000000e+01 9.5476758e-01 8.5647416e-01\n",
      "   1.0000000e+00]]], shape=(2, 1, 6), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# 多头注意力函数 测试\n",
    "\n",
    "batch_size = 2\n",
    "head_num = 2\n",
    "depth = 3\n",
    "d_model = 6\n",
    "\n",
    "\n",
    "def split_heads(x, batch_size):\n",
    "    # 分头, 将 \"头个数\" 的维度放到 seq_len 前面\n",
    "    \n",
    "    x = tf.reshape(x, (batch_size, -1, head_num, depth))\n",
    "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
    "\n",
    "\n",
    "querys = split_heads(test_q, batch_size) # (batch_size, num_heads, seq_len_q, depth)\n",
    "keys = split_heads(test_k, batch_size) # (batch_size, num_heads, seq_len_k, depth)\n",
    "values = split_heads(test_v, batch_size) # (batch_size, num_heads, seq_len_v, depth)\n",
    "print(\"\\n\")\n",
    "print(\"querys: \")\n",
    "print(querys)\n",
    "print(\"\\n\")\n",
    "\n",
    "mask = padding_mask(tf.constant([[1, 2, 0, 0], [3, 0, 1, 1]]))\n",
    "multi_head_test_out = scaled_dot_product_attention(querys, keys, values, mask) # (batch_size, num_heads, seq_len_q, depth)\n",
    "print(\"\\n\")\n",
    "print(\"multi_head_test_out: \")\n",
    "print(multi_head_test_out)\n",
    "\n",
    "\n",
    "# 把多头维度后移\n",
    "scaled_attention = tf.transpose(multi_head_test_out, [0, 2, 1, 3]) # (batch_size, seq_len_q, num_heads, depth)\n",
    "print(\"\\n\")\n",
    "print(\"scaled_attention: \")\n",
    "print(scaled_attention)\n",
    "\n",
    "# 合并多头\n",
    "concat_attention = tf.reshape(scaled_attention, (batch_size, -1, d_model))\n",
    "\n",
    "print(\"\\n\")\n",
    "print(\"concat_attention: \")\n",
    "print(concat_attention)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 构造 multi head attention 层\n",
    "\n",
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "\n",
    "        # d_model 必须可以正确分为各个头\n",
    "        assert d_model % num_heads == 0\n",
    "        \n",
    "        # 分头后的维度\n",
    "        self.depth = d_model // num_heads\n",
    "\n",
    "        self.wq = tf.keras.layers.Dense(d_model)\n",
    "        self.wk = tf.keras.layers.Dense(d_model)\n",
    "        self.wv = tf.keras.layers.Dense(d_model)\n",
    "\n",
    "        self.dense = tf.keras.layers.Dense(d_model)\n",
    "\n",
    "    def split_heads(self, x, batch_size):\n",
    "        # 分头, 将头个数的维度 放到 seq_len 前面\n",
    "        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
    "        return tf.transpose(x, perm=[0, 2, 1, 3])\n",
    "\n",
    "    def call(self, inputs):\n",
    "        q, k, v, mask = inputs\n",
    "        batch_size = tf.shape(q)[0]\n",
    "\n",
    "        # 分头前的前向网络，获取q、k、v语义\n",
    "        q = self.wq(q) # (batch_size, seq_len, d_model)\n",
    "        k = self.wk(k)\n",
    "        v = self.wv(v)\n",
    "\n",
    "        # 分头\n",
    "        q = self.split_heads(q, batch_size) # (batch_size, num_heads, seq_len_q, depth)\n",
    "        k = self.split_heads(k, batch_size) # (batch_size, num_heads, seq_len_k, depth)\n",
    "        v = self.split_heads(v, batch_size) # (batch_size, num_heads, seq_len_v, depth)\n",
    "        \n",
    "        # 通过缩放点积注意力层\n",
    "        scaled_attention = scaled_dot_product_attention(q, k, v, mask) # (batch_size, num_heads, seq_len_q, depth)\n",
    "        \n",
    "        # “多头维度” 后移\n",
    "        scaled_attention = tf.transpose(scaled_attention, [0, 2, 1, 3]) # (batch_size, seq_len_q, num_heads, depth)\n",
    "\n",
    "        # 合并 “多头维度”\n",
    "        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))\n",
    "\n",
    "        # 全连接层\n",
    "        output = self.dense(concat_attention)\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[[ 0.10577798 -0.25881276  0.41238633 ... -1.0608275  -0.49233449\n",
      "   -0.71212834]\n",
      "  [ 0.10273671 -0.25958157  0.4206068  ... -1.0604633  -0.48400968\n",
      "   -0.7194559 ]\n",
      "  [ 0.11328715 -0.2694013   0.41302037 ... -1.0594761  -0.4867704\n",
      "   -0.71812385]\n",
      "  ...\n",
      "  [ 0.1060189  -0.25882116  0.4104219  ... -1.0648378  -0.4930231\n",
      "   -0.7188878 ]\n",
      "  [ 0.10444206 -0.26108062  0.4155717  ... -1.0717115  -0.4878379\n",
      "   -0.7211157 ]\n",
      "  [ 0.10722721 -0.26775098  0.41619438 ... -1.0661967  -0.4871063\n",
      "   -0.71592176]]\n",
      "\n",
      " [[ 0.18043017 -0.22418204  0.5168557  ... -1.0312057  -0.4068474\n",
      "   -0.6909212 ]\n",
      "  [ 0.17852384 -0.22538428  0.519729   ... -1.0182251  -0.40997308\n",
      "   -0.68516874]\n",
      "  [ 0.17620301 -0.22790715  0.5144684  ... -1.0203692  -0.41394687\n",
      "   -0.68758833]\n",
      "  ...\n",
      "  [ 0.17873383 -0.22363052  0.5182944  ... -1.0251366  -0.41794103\n",
      "   -0.69461   ]\n",
      "  [ 0.17429072 -0.22895774  0.51766104 ... -1.0286089  -0.41128057\n",
      "   -0.6883719 ]\n",
      "  [ 0.18109483 -0.22909716  0.51299834 ... -1.0268948  -0.41131327\n",
      "   -0.6860363 ]]\n",
      "\n",
      " [[ 0.17741656 -0.28382164  0.4938854  ... -1.0981822  -0.35374826\n",
      "   -0.7242091 ]\n",
      "  [ 0.17866766 -0.2907327   0.4936837  ... -1.0978956  -0.34954894\n",
      "   -0.7177025 ]\n",
      "  [ 0.16877925 -0.29478204  0.49461415 ... -1.1045288  -0.35637265\n",
      "   -0.72271687]\n",
      "  ...\n",
      "  [ 0.18370318 -0.29095662  0.49587747 ... -1.0921474  -0.35413867\n",
      "   -0.72269255]\n",
      "  [ 0.18012899 -0.28966892  0.4991883  ... -1.0985599  -0.3531659\n",
      "   -0.72506595]\n",
      "  [ 0.17860287 -0.2914352   0.50017595 ... -1.0965195  -0.35573202\n",
      "   -0.723512  ]]\n",
      "\n",
      " [[ 0.1402139  -0.28023383  0.47633463 ... -1.0492234  -0.41984197\n",
      "   -0.7538934 ]\n",
      "  [ 0.15106416 -0.27861738  0.47906715 ... -1.0471029  -0.42565712\n",
      "   -0.7574363 ]\n",
      "  [ 0.14602816 -0.27822798  0.48509365 ... -1.0492668  -0.41652328\n",
      "   -0.75200677]\n",
      "  ...\n",
      "  [ 0.14039797 -0.27205113  0.4819404  ... -1.0485747  -0.42387348\n",
      "   -0.75484943]\n",
      "  [ 0.1399566  -0.2729967   0.48285967 ... -1.0483974  -0.42181724\n",
      "   -0.7529214 ]\n",
      "  [ 0.14408296 -0.2728199   0.480062   ... -1.0461028  -0.41725463\n",
      "   -0.7572432 ]]\n",
      "\n",
      " [[ 0.13372934 -0.24212328  0.44693744 ... -1.0488826  -0.4807039\n",
      "   -0.7557041 ]\n",
      "  [ 0.13070667 -0.2384171   0.45079052 ... -1.0501841  -0.48005706\n",
      "   -0.7525021 ]\n",
      "  [ 0.13536966 -0.2389047   0.44847754 ... -1.0555907  -0.48414958\n",
      "   -0.74958915]\n",
      "  ...\n",
      "  [ 0.13269931 -0.23876955  0.44201308 ... -1.0482553  -0.47982424\n",
      "   -0.7495153 ]\n",
      "  [ 0.13138556 -0.24047425  0.44516873 ... -1.0424248  -0.48641813\n",
      "   -0.75195503]\n",
      "  [ 0.12786096 -0.23880444  0.44373167 ... -1.0435371  -0.48390645\n",
      "   -0.7501143 ]]], shape=(5, 60, 512), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# 测试 MultiHeadAttention\n",
    "\n",
    "temp_mha = MultiHeadAttention(d_model=512, num_heads=8)\n",
    "input_data = tf.random.uniform((5, 60, 512))\n",
    "mask = None\n",
    "\n",
    "output = temp_mha([input_data, input_data, input_data, mask])\n",
    "print(output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "def point_wise_feed_forward_network(d_model, middle_units):\n",
    "    \n",
    "    return tf.keras.Sequential([\n",
    "        tf.keras.layers.Dense(middle_units, activation='relu'),\n",
    "        tf.keras.layers.Dense(d_model)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n# point_wise_feed_forward_network 测试\\n\\nsample_fnn = point_wise_feed_forward_network(512, 2048)\\ninput_data = tf.random.uniform((64, 50, 512))\\n\\npoint_wise_feed_forward_network = sample_fnn(input_data)\\nprint(point_wise_feed_forward_network.shape)\\n\\n'"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "# point_wise_feed_forward_network 测试\n",
    "\n",
    "sample_fnn = point_wise_feed_forward_network(512, 2048)\n",
    "input_data = tf.random.uniform((64, 50, 512))\n",
    "\n",
    "point_wise_feed_forward_network = sample_fnn(input_data)\n",
    "print(point_wise_feed_forward_network.shape)\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNormalization(tf.keras.layers.Layer):\n",
    "    def __init__(self, epsilon=1e-6, **kwargs):\n",
    "        self.eps = epsilon\n",
    "        super(LayerNormalization, self).__init__(**kwargs)\n",
    "        \n",
    "    def build(self, input_shape):\n",
    "        self.gamma = self.add_weight(name='gamma', shape=input_shape[-1:],\n",
    "                                     initializer=tf.ones_initializer(), trainable=True)\n",
    "        self.beta = self.add_weight(name='beta', shape=input_shape[-1:],\n",
    "                                    initializer=tf.zeros_initializer(), trainable=True)\n",
    "        super(LayerNormalization, self).build(input_shape)\n",
    "        \n",
    "    def call(self, x):\n",
    "        mean = tf.keras.backend.mean(x, axis=-1, keepdims=True)\n",
    "        std = tf.keras.backend.std(x, axis=-1, keepdims=True)\n",
    "        return self.gamma * (x - mean) / (std + self.eps) + self.beta\n",
    "    \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        return input_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, d_model, num_heads, middle_units, epsilon=1e-6, dropout_rate=0.1):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        \n",
    "        self.mha = MultiHeadAttention(d_model, num_heads)\n",
    "        self.ffn = point_wise_feed_forward_network(d_model, middle_units)\n",
    "        \n",
    "        self.layernorm1 = LayerNormalization()\n",
    "        self.layernorm2 = LayerNormalization()\n",
    "        \n",
    "        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n",
    "        \n",
    "    def call(self, inputs, mask, training):\n",
    "        # 多头注意力网络\n",
    "        att_output = self.mha([inputs, inputs, inputs, mask])\n",
    "        att_output = self.dropout1(att_output, training=training)\n",
    "        out1 = self.layernorm1(inputs + att_output)  # (batch_size, input_seq_len, d_model)\n",
    "        \n",
    "        # 前向网络\n",
    "        ffn_output = self.ffn(out1)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        out2 = self.layernorm2(out1 + ffn_output)   # (batch_size, input_seq_len, d_model)\n",
    "        \n",
    "        return out2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 43, 512)\n"
     ]
    }
   ],
   "source": [
    "# EncoderLayer 测试\n",
    "\n",
    "sample_encoder_layer = EncoderLayer(512, 8, 2048)\n",
    "input_data = tf.random.uniform((64, 43, 512))\n",
    "\n",
    "sample_encoder_layer_output = sample_encoder_layer(input_data, None, False)\n",
    "print(sample_encoder_layer_output.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(tf.keras.layers.Layer):\n",
    "    def __init__(self, n_layers, d_model, num_heads, middle_units,\n",
    "                max_seq_len, epsilon=1e-6, dropout_rate=0.1):\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        self.n_layers = n_layers\n",
    "        self.d_model = d_model\n",
    "        self.pos_embedding = PositionalEncoding(sequence_len=max_seq_len, embedding_dim=d_model)\n",
    "\n",
    "        self.encode_layer = [EncoderLayer(d_model=d_model, num_heads=num_heads, \n",
    "                                          middle_units=middle_units, \n",
    "                                          epsilon=epsilon, dropout_rate=dropout_rate)\n",
    "                            for _ in range(n_layers)]\n",
    "        \n",
    "    def call(self, inputs, mask, training):\n",
    "        emb = inputs\n",
    "        emb = self.pos_embedding(emb)\n",
    "        \n",
    "        for i in range(self.n_layers):\n",
    "            emb = self.encode_layer[i](emb, mask, training)\n",
    "\n",
    "        return emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(15, 60, 512)\n"
     ]
    }
   ],
   "source": [
    "# Encoder 测试\n",
    "\n",
    "sample_encoder = Encoder(2, 512, 8, 1024, 60)\n",
    "sample_encoder_output = sample_encoder(tf.random.uniform((15, 60, 512)),\n",
    "                                      None, False)\n",
    "\n",
    "print(sample_encoder_output.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data...\n",
      "25000 train sequences\n",
      "25000 test sequences\n",
      "x_train shape: (25000, 64)\n",
      "x_test shape: (25000, 64)\n"
     ]
    }
   ],
   "source": [
    "# 文本分类实验\n",
    "\n",
    "from tensorflow.keras.preprocessing import sequence\n",
    "from tensorflow.keras.datasets import imdb\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import SGD, Adam\n",
    "from tensorflow.keras.layers import *\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "\n",
    "# 1. 数据信息\n",
    "max_features = 20000\n",
    "maxlen = 64\n",
    "batch_size = 32\n",
    "\n",
    "\n",
    "print('Loading data...')\n",
    "(x_train, y_train), (x_test, y_test) = imdb.load_data(path=\"/notebook/Keras-Recommendation/scripts/multi_head_attention/imdb.npz\", \\\n",
    "                                                      num_words=max_features)\n",
    "y_train, y_test = pd.get_dummies(y_train), pd.get_dummies(y_test)\n",
    "print(len(x_train), 'train sequences')\n",
    "print(len(x_test), 'test sequences')\n",
    "\n",
    "\n",
    "x_train = sequence.pad_sequences(x_train, maxlen=maxlen)\n",
    "x_test = sequence.pad_sequences(x_test, maxlen=maxlen)\n",
    "\n",
    "print('x_train shape:', x_train.shape)\n",
    "print('x_test shape:', x_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "embeddings:\n",
      "Tensor(\"embedding_7/Identity:0\", shape=(None, 64, 128), dtype=float32)\n",
      "\n",
      "\n",
      "\n",
      "out_seq:\n",
      "Tensor(\"encoder_10/Identity:0\", shape=(None, 64, 128), dtype=float32)\n",
      "\n",
      "\n",
      "\n",
      "out_seq:\n",
      "Tensor(\"global_average_pooling1d_6/Identity:0\", shape=(None, 128), dtype=float32)\n",
      "Model: \"model_6\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_8 (InputLayer)            [(None, 64)]         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Equal_4 (TensorFlow [(None, 64)]         0           input_8[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_Cast_6 (TensorFlowO [(None, 64)]         0           tf_op_layer_Equal_4[0][0]        \n",
      "__________________________________________________________________________________________________\n",
      "embedding_7 (Embedding)         (None, 64, 128)      2560000     input_8[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "tf_op_layer_strided_slice_4 (Te [(None, 1, 1, 64)]   0           tf_op_layer_Cast_6[0][0]         \n",
      "__________________________________________________________________________________________________\n",
      "encoder_10 (Encoder)            (None, 64, 128)      264960      embedding_7[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "global_average_pooling1d_6 (Glo (None, 128)          0           encoder_10[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_60 (Dropout)            (None, 128)          0           global_average_pooling1d_6[0][0] \n",
      "__________________________________________________________________________________________________\n",
      "dropout_61 (Dropout)            (None, 128)          0           dropout_60[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dropout_62 (Dropout)            (None, 128)          0           dropout_61[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "dense_182 (Dense)               (None, 2)            258         dropout_62[0][0]                 \n",
      "==================================================================================================\n",
      "Total params: 2,825,218\n",
      "Trainable params: 2,825,218\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "None\n",
      "Train...\n",
      "Train on 25000 samples, validate on 25000 samples\n",
      "Epoch 1/10\n",
      "25000/25000 [==============================] - 125s 5ms/sample - loss: 0.7578 - accuracy: 0.5043 - val_loss: 0.6901 - val_accuracy: 0.5000\n",
      "Epoch 2/10\n",
      "25000/25000 [==============================] - 122s 5ms/sample - loss: 0.5051 - accuracy: 0.7265 - val_loss: 0.3716 - val_accuracy: 0.8341\n",
      "Epoch 3/10\n",
      "25000/25000 [==============================] - 123s 5ms/sample - loss: 0.3186 - accuracy: 0.8649 - val_loss: 0.3712 - val_accuracy: 0.8370\n",
      "Epoch 4/10\n",
      "25000/25000 [==============================] - 122s 5ms/sample - loss: 0.2462 - accuracy: 0.9030 - val_loss: 0.4216 - val_accuracy: 0.8328\n",
      "Epoch 5/10\n",
      "25000/25000 [==============================] - 122s 5ms/sample - loss: 0.1799 - accuracy: 0.9336 - val_loss: 0.4347 - val_accuracy: 0.8194\n",
      "Epoch 6/10\n",
      "25000/25000 [==============================] - 123s 5ms/sample - loss: 0.1168 - accuracy: 0.9600 - val_loss: 0.6673 - val_accuracy: 0.7985\n",
      "Epoch 7/10\n",
      "25000/25000 [==============================] - 119s 5ms/sample - loss: 0.0645 - accuracy: 0.9800 - val_loss: 0.7854 - val_accuracy: 0.8001\n",
      "Epoch 8/10\n",
      "25000/25000 [==============================] - 124s 5ms/sample - loss: 0.0320 - accuracy: 0.9910 - val_loss: 0.9486 - val_accuracy: 0.7840\n",
      "Epoch 9/10\n",
      "25000/25000 [==============================] - 122s 5ms/sample - loss: 0.0247 - accuracy: 0.9928 - val_loss: 1.2633 - val_accuracy: 0.7777\n",
      "Epoch 10/10\n",
      "25000/25000 [==============================] - 124s 5ms/sample - loss: 0.0176 - accuracy: 0.9952 - val_loss: 1.4019 - val_accuracy: 0.7814\n"
     ]
    }
   ],
   "source": [
    "# 2. 构造模型，及训练模型\n",
    "\n",
    "inputs = Input(shape=(64,), dtype='int32')\n",
    "embeddings = Embedding(max_features, 128)(inputs)\n",
    "\n",
    "print(\"\\n\"*2)\n",
    "print(\"embeddings:\")\n",
    "print(embeddings)\n",
    "\n",
    "mask_inputs = padding_mask(inputs)\n",
    "\n",
    "out_seq = Encoder(2, 128, 4, 256, maxlen)(embeddings, mask_inputs, False)\n",
    "\n",
    "print(\"\\n\"*2)\n",
    "print(\"out_seq:\")\n",
    "print(out_seq)\n",
    "\n",
    "out_seq = GlobalAveragePooling1D()(out_seq)\n",
    "\n",
    "print(\"\\n\"*2)\n",
    "print(\"out_seq:\")\n",
    "print(out_seq)\n",
    "\n",
    "out_seq = Dropout(0.3)(out_seq)\n",
    "outputs = Dense(64, activation='relu')(out_seq)\n",
    "\n",
    "out_seq = Dropout(0.3)(out_seq)\n",
    "outputs = Dense(16, activation='relu')(out_seq)\n",
    "\n",
    "out_seq = Dropout(0.3)(out_seq)\n",
    "outputs = Dense(2, activation='softmax')(out_seq)\n",
    "\n",
    "model = Model(inputs=inputs, outputs=outputs)\n",
    "print(model.summary())\n",
    "\n",
    "\n",
    "opt = Adam(lr=0.0002, decay=0.00001)\n",
    "loss = 'categorical_crossentropy'\n",
    "model.compile(loss=loss,\n",
    "             optimizer=opt,\n",
    "             metrics=['accuracy'])\n",
    "\n",
    "\n",
    "print('Train...')\n",
    "history = model.fit(x_train, y_train,\n",
    "         batch_size=batch_size,\n",
    "         epochs=10,\n",
    "         validation_data=(x_test, y_test))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
